diff --git a/arch/sandbox/include/asm/malloc.h b/arch/sandbox/include/asm/malloc.h
index a1467b5..8aaaa9c 100644
--- a/arch/sandbox/include/asm/malloc.h
+++ b/arch/sandbox/include/asm/malloc.h
@@ -6,6 +6,7 @@
  */
 
 #ifndef __ASM_MALLOC_H
+#define __ASM_MALLOC_H
 
 void *malloc(size_t size);
 void free(void *ptr);
diff --git a/common/dlmalloc.c b/common/dlmalloc.c
index f48cd2a..41c7230 100644
--- a/common/dlmalloc.c
+++ b/common/dlmalloc.c
@@ -596,6 +596,9 @@
 ulong mem_malloc_end = 0;
 ulong mem_malloc_brk = 0;
 
+static bool malloc_testing;	/* enable test mode */
+static int malloc_max_allocs;	/* return NULL after this many calls to malloc() */
+
 void *sbrk(ptrdiff_t increment)
 {
 	ulong old = mem_malloc_brk;
@@ -1307,6 +1310,11 @@
 		return malloc_simple(bytes);
 #endif
 
+  if (CONFIG_IS_ENABLED(UNIT_TEST) && malloc_testing) {
+    if (--malloc_max_allocs < 0)
+      return NULL;
+  }
+
   /* check if mem_malloc_init() was run */
   if ((mem_malloc_start == 0) && (mem_malloc_end == 0)) {
     /* not initialized yet */
@@ -2470,6 +2478,17 @@
 	return 0;
 }
 
+void malloc_enable_testing(int max_allocs)
+{
+	malloc_testing = true;
+	malloc_max_allocs = max_allocs;
+}
+
+void malloc_disable_testing(void)
+{
+	malloc_testing = false;
+}
+
 /*
 
 History:
diff --git a/include/malloc.h b/include/malloc.h
index e8c8b25..161ccbd 100644
--- a/include/malloc.h
+++ b/include/malloc.h
@@ -883,6 +883,18 @@
 
 void malloc_simple_info(void);
 
+/**
+ * malloc_enable_testing() - Put malloc() into test mode
+ *
+ * This only works if UNIT_TESTING is enabled
+ *
+ * @max_allocs: return -ENOMEM after max_allocs calls to malloc()
+ */
+void malloc_enable_testing(int max_allocs);
+
+/** malloc_disable_testing() - Put malloc() into normal mode */
+void malloc_disable_testing(void);
+
 #if CONFIG_IS_ENABLED(SYS_MALLOC_SIMPLE)
 #define malloc malloc_simple
 #define realloc realloc_simple
diff --git a/test/test-main.c b/test/test-main.c
index 90a324b..7a38716 100644
--- a/test/test-main.c
+++ b/test/test-main.c
@@ -47,6 +47,7 @@
 	uts->force_fail_alloc = false;
 	uts->skip_post_probe = false;
 	gd->dm_root = NULL;
+	malloc_disable_testing();
 	if (CONFIG_IS_ENABLED(UT_DM) && !CONFIG_IS_ENABLED(OF_PLATDATA))
 		memset(dm_testdrv_op_count, '\0', sizeof(dm_testdrv_op_count));
 	arch_reset_for_test();